from PIL import Image
import os
import re
import tkinter as tk
from tkinter import filedialog
from pathlib import Path

# 自然排序算法（支持中文+数字混合排序）
def numerical_sort(value):
    parts = re.split(r'(\d+)', str(value))
    parts[1::2] = map(int, parts[1::2])
    return parts

# 弹窗式路径选择
def select_paths_with_prompt():
    root = tk.Tk()
    root.withdraw()
    
    # 第一步：选择源目录
    source_folder = filedialog.askdirectory(
        title="第一步：请选择包含图片子目录的文件夹",
        mustexist=True
    )
    if not source_folder:
        print("未选择源目录，程序退出")
        exit()

    # 第二步：选择保存路径
    default_name = f"{Path(source_folder).name}_合并结果.pdf"
    save_path = filedialog.asksaveasfilename(
        title="第二步：请指定PDF保存路径和文件名",
        initialdir=Path(source_folder).parent,
        initialfile=default_name,
        defaultextension=".pdf",
        filetypes=[("PDF文件", "*.pdf")],
        confirmoverwrite=True  # 防止覆盖确认
    )
    return source_folder, save_path

# 处理所有子目录
def process_subfolders(source_folder, output_pdf):
    # 按自然顺序排序子目录
    subfolders = sorted(
        [f for f in Path(source_folder).iterdir() if f.is_dir()],
        key=lambda x: numerical_sort(x.name)
    )
    
    # 分批次处理图片（内存优化）
    first_page = None
    try:
        for folder in subfolders:
            # 按自然顺序获取图片文件
            image_files = sorted(
                folder.glob("*.*"),
                key=lambda x: numerical_sort(x.stem)
            )
            image_files = [
                f for f in image_files 
                if f.suffix.lower() in ('.jpg', '.jpeg', '.png')
            ]
            
            # 处理当前文件夹的图片
            for img_path in image_files:
                with Image.open(img_path) as img:
                    if img.mode != "RGB":
                        img = img.convert("RGB")
                    
                    # 处理第一张图片
                    if not first_page:
                        first_page = img.copy()
                        first_page.save(output_pdf, "PDF", save_all=True)
                        print(f"初始化PDF：{output_pdf}")
                    # 追加后续图片
                    else:
                        img.save(output_pdf, "PDF", save_all=True, append=True)
                        
                    print(f"已合并：{folder.name}/{img_path.name}")
                    
    except Exception as e:
        print(f"处理中断！错误原因：{str(e)}")
        # 删除未完成的PDF文件
        if Path(output_pdf).exists():
            Path(output_pdf).unlink()
        exit()
            
    return output_pdf

if __name__ == "__main__":
    # 选择目录并执行合并
    source_folder, save_path = select_paths_with_prompt()
    if not save_path:
        print("未选择保存路径，程序退出")
        exit()
        
    result_pdf = process_subfolders(source_folder, save_path)
    print(f"合并完成！PDF文件保存至：{result_pdf}")